package club.babyf.nativeai.utils;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.regex.Pattern;

/**
 * TextSplitter.
 *
 * @author zhanglei.
 * @date 2025/3/17 11:36.
 * @description 智能分块工具.
 */
public class TextSplitter {

    private static final String ABBREVIATIONS = "(?i)\\b(Mr|Mrs|Ms|Dr|Prof|Rev|Jan|Feb|Aug|Sept|Oct|Nov|Dec|Ph\\.?D|U\\.?S|etc)\\.";

    private static final Pattern SENTENCE_PATTERN = Pattern.compile(
            "(?<!" + ABBREVIATIONS + ")(?<=[.!?。！？])\\s+(?=[\\p{Lu}\\p{IsHan}])"
    );

    public static List<String> splitIntoSensibleChunks(String text, int maxChunkSize, int overlap) {
        String[] sentences = SENTENCE_PATTERN.split(text);
        List<String> chunks = new ArrayList<>();
        int start = 0;
        while (start < sentences.length) {
            int optimalSize = calculateOptimalSize(sentences, start, maxChunkSize);
            if (optimalSize == 0) {
                // 确保至少有一个句子
                optimalSize = 1;
            }
            int end = Math.min(start + optimalSize, sentences.length);
            String chunk = String.join(" ", Arrays.copyOfRange(sentences, start, end));
            chunks.add(chunk);
            // 防止 start 变负数
            start = Math.max(end - overlap, start + 1);
        }
        return chunks;
    }

    private static int calculateOptimalSize(String[] sentences, int start, int maxSize) {
        int currentLength = 0;
        int lastValidIndex = start;
        for (int i = start; i < sentences.length; i++) {
            currentLength += sentences[i].length();
            if (currentLength > maxSize) break;
            lastValidIndex = i;
        }
        // 确保不切割句子
        return lastValidIndex - start + 1;
    }

}
